#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
RedFire 数据质量检查组件
提供完整性、准确性、一致性、时效性等多维度的数据质量检查
"""

import asyncio
import logging
import pandas as pd
from datetime import datetime, timedelta
from typing import Dict, List, Any, Optional, Tuple, Union
from dataclasses import dataclass, field
from enum import Enum
import json
import re

# 配置日志
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)


class QualityCheckType(Enum):
    """数据质量检查类型"""
    COMPLETENESS = "completeness"      # 完整性
    ACCURACY = "accuracy"              # 准确性
    CONSISTENCY = "consistency"        # 一致性
    TIMELINESS = "timeliness"          # 时效性
    VALIDITY = "validity"              # 有效性
    UNIQUENESS = "uniqueness"          # 唯一性


class QualityLevel(Enum):
    """数据质量等级"""
    EXCELLENT = "excellent"  # 优秀 (95-100%)
    GOOD = "good"           # 良好 (85-95%)
    FAIR = "fair"           # 一般 (70-85%)
    POOR = "poor"           # 较差 (0-70%)


@dataclass
class QualityRule:
    """数据质量规则"""
    rule_id: str
    name: str
    description: str
    check_type: QualityCheckType
    table_name: str
    column_name: Optional[str] = None
    rule_config: Dict[str, Any] = field(default_factory=dict)
    threshold: float = 0.95  # 质量阈值
    is_active: bool = True


@dataclass
class QualityCheckResult:
    """数据质量检查结果"""
    rule_id: str
    rule_name: str
    check_type: QualityCheckType
    table_name: str
    column_name: Optional[str]
    score: float  # 0-1之间的得分
    level: QualityLevel
    pass_threshold: bool
    total_records: int
    valid_records: int
    invalid_records: int
    details: Dict[str, Any] = field(default_factory=dict)
    error_samples: List[Dict[str, Any]] = field(default_factory=list)
    checked_at: datetime = field(default_factory=datetime.now)


@dataclass
class QualityReport:
    """数据质量报告"""
    report_id: str
    table_name: str
    overall_score: float
    overall_level: QualityLevel
    check_results: List[QualityCheckResult]
    summary: Dict[str, Any] = field(default_factory=dict)
    recommendations: List[str] = field(default_factory=list)
    generated_at: datetime = field(default_factory=datetime.now)


class DataQualityChecker:
    """数据质量检查器"""
    
    def __init__(self, db_config: Dict[str, Any]):
        self.db_config = db_config
        self.connection = None
        self.quality_rules = []
        self._load_default_rules()
    
    async def connect(self):
        """建立数据库连接"""
        if self.db_config.get('type') == 'postgresql':
            import asyncpg
            self.connection = await asyncpg.connect(
                host=self.db_config['host'],
                port=self.db_config['port'],
                database=self.db_config['database'],
                user=self.db_config['username'],
                password=self.db_config['password']
            )
        else:
            raise ValueError(f"不支持的数据库类型: {self.db_config.get('type')}")
    
    async def close(self):
        """关闭数据库连接"""
        if self.connection:
            await self.connection.close()
    
    def _load_default_rules(self):
        """加载默认质量规则"""
        # Tick数据完整性规则
        self.quality_rules.extend([
            QualityRule(
                rule_id="tick_completeness_symbol",
                name="Tick数据品种代码完整性",
                description="检查Tick数据中品种代码字段的完整性",
                check_type=QualityCheckType.COMPLETENESS,
                table_name="dbtickdata",
                column_name="symbol",
                threshold=0.99
            ),
            QualityRule(
                rule_id="tick_completeness_price",
                name="Tick数据价格完整性",
                description="检查Tick数据中价格字段的完整性",
                check_type=QualityCheckType.COMPLETENESS,
                table_name="dbtickdata",
                column_name="last_price",
                threshold=0.95
            ),
            QualityRule(
                rule_id="tick_validity_price_range",
                name="Tick数据价格有效性",
                description="检查Tick数据价格是否在合理范围内",
                check_type=QualityCheckType.VALIDITY,
                table_name="dbtickdata",
                column_name="last_price",
                rule_config={
                    "min_value": 0,
                    "max_value": 100000
                },
                threshold=0.98
            ),
            QualityRule(
                rule_id="tick_timeliness_delay",
                name="Tick数据时效性",
                description="检查Tick数据的时间延迟",
                check_type=QualityCheckType.TIMELINESS,
                table_name="dbtickdata",
                column_name="datetime",
                rule_config={
                    "max_delay_minutes": 30
                },
                threshold=0.90
            ),
            QualityRule(
                rule_id="tick_uniqueness_primary",
                name="Tick数据唯一性",
                description="检查Tick数据的主键唯一性",
                check_type=QualityCheckType.UNIQUENESS,
                table_name="dbtickdata",
                rule_config={
                    "unique_columns": ["symbol", "exchange", "datetime"]
                },
                threshold=1.0
            )
        ])
        
        # K线数据质量规则
        self.quality_rules.extend([
            QualityRule(
                rule_id="bar_completeness_ohlc",
                name="K线数据OHLC完整性",
                description="检查K线数据OHLC字段的完整性",
                check_type=QualityCheckType.COMPLETENESS,
                table_name="dbbardata",
                rule_config={
                    "required_columns": ["open_price", "high_price", "low_price", "close_price"]
                },
                threshold=0.99
            ),
            QualityRule(
                rule_id="bar_consistency_ohlc",
                name="K线数据OHLC一致性",
                description="检查K线数据OHLC价格的逻辑一致性",
                check_type=QualityCheckType.CONSISTENCY,
                table_name="dbbardata",
                rule_config={
                    "check_ohlc_logic": True
                },
                threshold=0.95
            ),
            QualityRule(
                rule_id="bar_accuracy_volume",
                name="K线数据成交量准确性",
                description="检查K线数据成交量的准确性",
                check_type=QualityCheckType.ACCURACY,
                table_name="dbbardata",
                column_name="volume",
                rule_config={
                    "min_volume": 0,
                    "check_negative": True
                },
                threshold=0.98
            )
        ])
    
    def add_quality_rule(self, rule: QualityRule):
        """添加质量规则"""
        self.quality_rules.append(rule)
        logger.info(f"添加质量规则: {rule.name}")
    
    def remove_quality_rule(self, rule_id: str):
        """移除质量规则"""
        self.quality_rules = [r for r in self.quality_rules if r.rule_id != rule_id]
        logger.info(f"移除质量规则: {rule_id}")
    
    async def check_table_quality(self, table_name: str, 
                                 rules: List[str] = None) -> QualityReport:
        """检查指定表的数据质量"""
        await self.connect()
        
        try:
            # 过滤适用的规则
            applicable_rules = [
                rule for rule in self.quality_rules
                if rule.table_name == table_name and rule.is_active
                and (rules is None or rule.rule_id in rules)
            ]
            
            if not applicable_rules:
                logger.warning(f"表 {table_name} 没有适用的质量规则")
                return QualityReport(
                    report_id=f"{table_name}_{int(datetime.now().timestamp())}",
                    table_name=table_name,
                    overall_score=0.0,
                    overall_level=QualityLevel.POOR,
                    check_results=[]
                )
            
            # 执行质量检查
            check_results = []
            for rule in applicable_rules:
                try:
                    result = await self._execute_quality_check(rule)
                    check_results.append(result)
                    logger.info(f"质量检查完成: {rule.name}, 得分: {result.score:.2%}")
                except Exception as e:
                    logger.error(f"质量检查失败: {rule.name}, 错误: {e}")
            
            # 生成综合报告
            report = self._generate_quality_report(table_name, check_results)
            
            # 保存检查结果
            await self._save_quality_results(check_results)
            
            return report
            
        finally:
            await self.close()
    
    async def _execute_quality_check(self, rule: QualityRule) -> QualityCheckResult:
        """执行单个质量检查"""
        if rule.check_type == QualityCheckType.COMPLETENESS:
            return await self._check_completeness(rule)
        elif rule.check_type == QualityCheckType.ACCURACY:
            return await self._check_accuracy(rule)
        elif rule.check_type == QualityCheckType.CONSISTENCY:
            return await self._check_consistency(rule)
        elif rule.check_type == QualityCheckType.TIMELINESS:
            return await self._check_timeliness(rule)
        elif rule.check_type == QualityCheckType.VALIDITY:
            return await self._check_validity(rule)
        elif rule.check_type == QualityCheckType.UNIQUENESS:
            return await self._check_uniqueness(rule)
        else:
            raise ValueError(f"不支持的检查类型: {rule.check_type}")
    
    async def _check_completeness(self, rule: QualityRule) -> QualityCheckResult:
        """检查完整性"""
        if rule.column_name:
            # 单列完整性检查
            sql = f"""
            SELECT 
                COUNT(*) as total_count,
                COUNT({rule.column_name}) as valid_count,
                COUNT(*) - COUNT({rule.column_name}) as null_count
            FROM {rule.table_name}
            WHERE created_at >= NOW() - INTERVAL '1 day'
            """
        else:
            # 多列完整性检查
            required_columns = rule.rule_config.get('required_columns', [])
            column_checks = []
            for col in required_columns:
                column_checks.append(f"COUNT({col}) as {col}_count")
            
            sql = f"""
            SELECT 
                COUNT(*) as total_count,
                {', '.join(column_checks)}
            FROM {rule.table_name}
            WHERE created_at >= NOW() - INTERVAL '1 day'
            """
        
        result = await self.connection.fetchrow(sql)
        
        if rule.column_name:
            total_records = result['total_count']
            valid_records = result['valid_count']
            invalid_records = result['null_count']
        else:
            total_records = result['total_count']
            valid_records = min([result[f"{col}_count"] for col in required_columns])
            invalid_records = total_records - valid_records
        
        score = valid_records / total_records if total_records > 0 else 0
        level = self._calculate_quality_level(score)
        
        return QualityCheckResult(
            rule_id=rule.rule_id,
            rule_name=rule.name,
            check_type=rule.check_type,
            table_name=rule.table_name,
            column_name=rule.column_name,
            score=score,
            level=level,
            pass_threshold=score >= rule.threshold,
            total_records=total_records,
            valid_records=valid_records,
            invalid_records=invalid_records,
            details={
                'completeness_rate': score,
                'null_count': invalid_records
            }
        )
    
    async def _check_accuracy(self, rule: QualityRule) -> QualityCheckResult:
        """检查准确性"""
        config = rule.rule_config
        
        conditions = []
        if 'min_value' in config:
            conditions.append(f"{rule.column_name} >= {config['min_value']}")
        if 'max_value' in config:
            conditions.append(f"{rule.column_name} <= {config['max_value']}")
        if config.get('check_negative', False):
            conditions.append(f"{rule.column_name} >= 0")
        
        where_clause = " AND ".join(conditions) if conditions else "1=1"
        
        sql = f"""
        SELECT 
            COUNT(*) as total_count,
            COUNT(CASE WHEN {where_clause} THEN 1 END) as valid_count
        FROM {rule.table_name}
        WHERE {rule.column_name} IS NOT NULL
        AND created_at >= NOW() - INTERVAL '1 day'
        """
        
        result = await self.connection.fetchrow(sql)
        
        total_records = result['total_count']
        valid_records = result['valid_count']
        invalid_records = total_records - valid_records
        
        score = valid_records / total_records if total_records > 0 else 0
        level = self._calculate_quality_level(score)
        
        return QualityCheckResult(
            rule_id=rule.rule_id,
            rule_name=rule.name,
            check_type=rule.check_type,
            table_name=rule.table_name,
            column_name=rule.column_name,
            score=score,
            level=level,
            pass_threshold=score >= rule.threshold,
            total_records=total_records,
            valid_records=valid_records,
            invalid_records=invalid_records,
            details={
                'accuracy_rate': score,
                'invalid_count': invalid_records,
                'check_conditions': conditions
            }
        )
    
    async def _check_consistency(self, rule: QualityRule) -> QualityCheckResult:
        """检查一致性"""
        config = rule.rule_config
        
        if config.get('check_ohlc_logic', False):
            # 检查OHLC价格逻辑一致性
            sql = f"""
            SELECT 
                COUNT(*) as total_count,
                COUNT(CASE 
                    WHEN high_price >= low_price 
                    AND high_price >= open_price 
                    AND high_price >= close_price
                    AND low_price <= open_price 
                    AND low_price <= close_price
                    THEN 1 
                END) as consistent_count
            FROM {rule.table_name}
            WHERE open_price > 0 AND high_price > 0 AND low_price > 0 AND close_price > 0
            AND created_at >= NOW() - INTERVAL '1 day'
            """
            
            result = await self.connection.fetchrow(sql)
            
            total_records = result['total_count']
            valid_records = result['consistent_count']
            invalid_records = total_records - valid_records
            
            details = {
                'consistency_type': 'ohlc_logic',
                'consistency_rate': valid_records / total_records if total_records > 0 else 0
            }
        else:
            # 其他一致性检查
            total_records = 0
            valid_records = 0
            invalid_records = 0
            details = {}
        
        score = valid_records / total_records if total_records > 0 else 0
        level = self._calculate_quality_level(score)
        
        return QualityCheckResult(
            rule_id=rule.rule_id,
            rule_name=rule.name,
            check_type=rule.check_type,
            table_name=rule.table_name,
            column_name=rule.column_name,
            score=score,
            level=level,
            pass_threshold=score >= rule.threshold,
            total_records=total_records,
            valid_records=valid_records,
            invalid_records=invalid_records,
            details=details
        )
    
    async def _check_timeliness(self, rule: QualityRule) -> QualityCheckResult:
        """检查时效性"""
        config = rule.rule_config
        max_delay_minutes = config.get('max_delay_minutes', 30)
        
        sql = f"""
        SELECT 
            COUNT(*) as total_count,
            COUNT(CASE 
                WHEN {rule.column_name} >= NOW() - INTERVAL '{max_delay_minutes} minutes'
                THEN 1 
            END) as timely_count
        FROM {rule.table_name}
        WHERE {rule.column_name} IS NOT NULL
        AND created_at >= NOW() - INTERVAL '1 day'
        """
        
        result = await self.connection.fetchrow(sql)
        
        total_records = result['total_count']
        valid_records = result['timely_count']
        invalid_records = total_records - valid_records
        
        score = valid_records / total_records if total_records > 0 else 0
        level = self._calculate_quality_level(score)
        
        return QualityCheckResult(
            rule_id=rule.rule_id,
            rule_name=rule.name,
            check_type=rule.check_type,
            table_name=rule.table_name,
            column_name=rule.column_name,
            score=score,
            level=level,
            pass_threshold=score >= rule.threshold,
            total_records=total_records,
            valid_records=valid_records,
            invalid_records=invalid_records,
            details={
                'timeliness_rate': score,
                'max_delay_minutes': max_delay_minutes,
                'delayed_count': invalid_records
            }
        )
    
    async def _check_validity(self, rule: QualityRule) -> QualityCheckResult:
        """检查有效性"""
        config = rule.rule_config
        
        conditions = []
        if 'min_value' in config:
            conditions.append(f"{rule.column_name} >= {config['min_value']}")
        if 'max_value' in config:
            conditions.append(f"{rule.column_name} <= {config['max_value']}")
        if 'regex_pattern' in config:
            conditions.append(f"{rule.column_name} ~ '{config['regex_pattern']}'")
        
        where_clause = " AND ".join(conditions) if conditions else "1=1"
        
        sql = f"""
        SELECT 
            COUNT(*) as total_count,
            COUNT(CASE WHEN {where_clause} THEN 1 END) as valid_count
        FROM {rule.table_name}
        WHERE {rule.column_name} IS NOT NULL
        AND created_at >= NOW() - INTERVAL '1 day'
        """
        
        result = await self.connection.fetchrow(sql)
        
        total_records = result['total_count']
        valid_records = result['valid_count']
        invalid_records = total_records - valid_records
        
        score = valid_records / total_records if total_records > 0 else 0
        level = self._calculate_quality_level(score)
        
        return QualityCheckResult(
            rule_id=rule.rule_id,
            rule_name=rule.name,
            check_type=rule.check_type,
            table_name=rule.table_name,
            column_name=rule.column_name,
            score=score,
            level=level,
            pass_threshold=score >= rule.threshold,
            total_records=total_records,
            valid_records=valid_records,
            invalid_records=invalid_records,
            details={
                'validity_rate': score,
                'invalid_count': invalid_records,
                'validation_rules': conditions
            }
        )
    
    async def _check_uniqueness(self, rule: QualityRule) -> QualityCheckResult:
        """检查唯一性"""
        config = rule.rule_config
        unique_columns = config.get('unique_columns', [])
        
        if not unique_columns:
            raise ValueError("唯一性检查需要指定unique_columns")
        
        columns_str = ", ".join(unique_columns)
        
        sql = f"""
        SELECT 
            COUNT(*) as total_count,
            COUNT(DISTINCT ({columns_str})) as unique_count
        FROM {rule.table_name}
        WHERE created_at >= NOW() - INTERVAL '1 day'
        """
        
        result = await self.connection.fetchrow(sql)
        
        total_records = result['total_count']
        unique_records = result['unique_count']
        duplicate_records = total_records - unique_records
        
        score = unique_records / total_records if total_records > 0 else 0
        level = self._calculate_quality_level(score)
        
        return QualityCheckResult(
            rule_id=rule.rule_id,
            rule_name=rule.name,
            check_type=rule.check_type,
            table_name=rule.table_name,
            column_name=None,
            score=score,
            level=level,
            pass_threshold=score >= rule.threshold,
            total_records=total_records,
            valid_records=unique_records,
            invalid_records=duplicate_records,
            details={
                'uniqueness_rate': score,
                'duplicate_count': duplicate_records,
                'unique_columns': unique_columns
            }
        )
    
    def _calculate_quality_level(self, score: float) -> QualityLevel:
        """计算质量等级"""
        if score >= 0.95:
            return QualityLevel.EXCELLENT
        elif score >= 0.85:
            return QualityLevel.GOOD
        elif score >= 0.70:
            return QualityLevel.FAIR
        else:
            return QualityLevel.POOR
    
    def _generate_quality_report(self, table_name: str, 
                               check_results: List[QualityCheckResult]) -> QualityReport:
        """生成质量报告"""
        if not check_results:
            return QualityReport(
                report_id=f"{table_name}_{int(datetime.now().timestamp())}",
                table_name=table_name,
                overall_score=0.0,
                overall_level=QualityLevel.POOR,
                check_results=[]
            )
        
        # 计算综合得分
        overall_score = sum(result.score for result in check_results) / len(check_results)
        overall_level = self._calculate_quality_level(overall_score)
        
        # 生成摘要
        summary = {
            'total_checks': len(check_results),
            'passed_checks': sum(1 for r in check_results if r.pass_threshold),
            'failed_checks': sum(1 for r in check_results if not r.pass_threshold),
            'check_types': list(set(r.check_type.value for r in check_results)),
            'avg_score_by_type': {}
        }
        
        # 按检查类型计算平均得分
        for check_type in QualityCheckType:
            type_results = [r for r in check_results if r.check_type == check_type]
            if type_results:
                avg_score = sum(r.score for r in type_results) / len(type_results)
                summary['avg_score_by_type'][check_type.value] = avg_score
        
        # 生成建议
        recommendations = self._generate_recommendations(check_results)
        
        return QualityReport(
            report_id=f"{table_name}_{int(datetime.now().timestamp())}",
            table_name=table_name,
            overall_score=overall_score,
            overall_level=overall_level,
            check_results=check_results,
            summary=summary,
            recommendations=recommendations
        )
    
    def _generate_recommendations(self, check_results: List[QualityCheckResult]) -> List[str]:
        """生成改进建议"""
        recommendations = []
        
        # 检查失败的规则
        failed_results = [r for r in check_results if not r.pass_threshold]
        
        for result in failed_results:
            if result.check_type == QualityCheckType.COMPLETENESS:
                recommendations.append(
                    f"改进 {result.table_name}.{result.column_name or '多字段'} 的数据完整性，"
                    f"当前完整率 {result.score:.1%}"
                )
            elif result.check_type == QualityCheckType.ACCURACY:
                recommendations.append(
                    f"检查 {result.table_name}.{result.column_name} 的数据准确性，"
                    f"发现 {result.invalid_records} 条异常数据"
                )
            elif result.check_type == QualityCheckType.TIMELINESS:
                recommendations.append(
                    f"优化 {result.table_name} 的数据时效性，"
                    f"当前有 {result.invalid_records} 条延迟数据"
                )
            elif result.check_type == QualityCheckType.UNIQUENESS:
                recommendations.append(
                    f"解决 {result.table_name} 的数据重复问题，"
                    f"发现 {result.invalid_records} 条重复记录"
                )
        
        # 总体建议
        if len(failed_results) > len(check_results) * 0.5:
            recommendations.append("建议全面审查数据采集和处理流程")
        
        return recommendations
    
    async def _save_quality_results(self, check_results: List[QualityCheckResult]):
        """保存质量检查结果"""
        try:
            for result in check_results:
                sql = """
                INSERT INTO data_quality_log 
                (source_name, table_name, check_type, check_result, details)
                VALUES ($1, $2, $3, $4, $5)
                """
                
                check_result_status = 'PASS' if result.pass_threshold else 'FAIL'
                details = {
                    'rule_id': result.rule_id,
                    'score': result.score,
                    'level': result.level.value,
                    'total_records': result.total_records,
                    'valid_records': result.valid_records,
                    'invalid_records': result.invalid_records,
                    'details': result.details
                }
                
                await self.connection.execute(
                    sql,
                    'data_quality_checker',
                    result.table_name,
                    result.check_type.value,
                    check_result_status,
                    json.dumps(details, default=str)
                )
                
        except Exception as e:
            logger.error(f"保存质量检查结果失败: {e}")


# 使用示例
async def example_quality_check():
    """数据质量检查示例"""
    # 数据库配置
    db_config = {
        'type': 'postgresql',
        'host': 'localhost',
        'port': 5432,
        'database': 'redfire_stage2',
        'username': 'redfire_prod',
        'password': 'your_password'
    }
    
    # 创建质量检查器
    checker = DataQualityChecker(db_config)
    
    # 检查Tick数据质量
    tick_report = await checker.check_table_quality('dbtickdata')
    
    print("=== Tick数据质量报告 ===")
    print(f"综合得分: {tick_report.overall_score:.1%}")
    print(f"质量等级: {tick_report.overall_level.value}")
    print(f"检查项数: {len(tick_report.check_results)}")
    
    for result in tick_report.check_results:
        status = "✅" if result.pass_threshold else "❌"
        print(f"{status} {result.rule_name}: {result.score:.1%}")
    
    if tick_report.recommendations:
        print("\n改进建议:")
        for rec in tick_report.recommendations:
            print(f"- {rec}")


if __name__ == "__main__":
    asyncio.run(example_quality_check())
